'''
Created on Jan 16 2012

@author: oabalbin

This script takes
input: a list of snps positions by sample
output a file like CHR  POS    SP1   SP2   SP3
                    1   345    A/T   A/T   A/C
'''
import os
import sys
import subprocess
import numpy as np
from optparse import OptionParser
from collections import deque, defaultdict

                
def read_positions(ifile):
    '''
    '''
    ifile = open(ifile)
    for l in ifile:
        if l.startswith('#'):
            continue
        fields=l.strip('\n').split(',')
        fields.pop()
    
    ifile.close()
    return fields


def isec_arrays(filea,fileb):
    '''
    '''
    m1=set(read_positions(filea))
    m2=set(read_positions(fileb))
    #print sys.getsizeof(m1), sys.getsizeof(m2)
    isec=m1.intersection(m2)        
    return isec

def isec_snvs_annot_union(snvs_union, snvs_annot_dbs):
    '''
    This functions obtains all query snvs that are 
    common to a particular database, ex: dbSNP
    '''
    pm = defaultdict()
    for db, snv_db in snvs_annot_dbs.iteritems():
        pm[db] = set(snvs_union).intersection(set(snv_db))
           
    return pm

def union_snvs_dict(union_snvs):
    '''
    It gives an index to each snp. 
    So it can be used later to find arrays
    '''
    pm = defaultdict()
    for i, snv in enumerate(union_snvs):
        pm[snv]=i
    return pm


def read_patient_files(file):
    ifile=open(file)
    pat=defaultdict()
    i=0
    for l in ifile:
        pat[i]=l.strip('\n').split('\t')
        i+=1
    ifile.close()
    return pat


if __name__ == '__main__':
    '''
    Write to an array the positions in a particualr vcf.
    
    '''
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-f", "--pos_files", action="append", type="str",
                            dest="pos_files",
                            help="others_vcf_files. Pass vcf files without .gz extention even if the file .gz is there")
    
    optionparser.add_option("-p", "--patient_files", type="str", dest="patient_files",
                            help="others_vcf_files. Pass vcf files without .gz extention even if the file .gz is there")


    optionparser.add_option("-o", "--out_file", type="str",dest="ofile",
                            help="others_vcf_files. Pass vcf files without .gz extention even if the file .gz is there")
    
    optionparser.add_option("-b", "--base", dest="base", action="store_true", default=False,
                            help="determine if the ouput file is position or pos.base")
    optionparser.add_option("-g", "--geno", dest="geno", action="store_true", default=False,
                            help="determine if the ouput file is position or pos.base.geno")

    
    
    (options, args) = optionparser.parse_args()
    
    
    
    snvs_union = set()
    
    pfiles=defaultdict()
    snvs_annot_dbs = defaultdict()
    snvs_singletons = defaultdict()

    patients_files = read_patient_files(options.patient_files)
    # Read pos.base.geno files
    First=True
    datadir='/exds/users/oabalbin/projects/fmpn/data/gatk_calls'
    
    for i, thfile in patients_files.iteritems():   
    #for ft in options.pos_files:
        # Generate a positions files for query vcf.
        ft=thfile[2]
        if ft!='':
            print os.path.join(datadir,ft)
            thsnvs = read_positions(os.path.join(datadir,ft))
            #snvs_annot_dbs[ft]=thsnvs
            snvs_union.update(set(thsnvs))
            # isec set
            if First:
                isec_set=set(thsnvs)
                First=False
            else:
                isec_set.intersection_update( set(thsnvs) )
            
            # snvs present only in this samples
            snvs_singletons[thfile[0]]= set(thsnvs).difference(isec_set)
        
    diff_snvs = snvs_union.difference(isec_set)
    '''
    chr14@39272493|C>A|0/1
    '''

    # For the snvs shared across all individuals, isec set
    nrows=len(isec_set)+len(diff_snvs)
    ncols=1
    mytype=[('CHRM',int),('POS',int),('IBD',int),('MUT','S3'),('GENO','S3'),('SAMPLE',int)]
    mytype_ibd=[('CHRM',int),('POS',int),('IBD',int)]
  
    #location_arr=np.empty(( nrows, ncols),dtype=mytype)
    all_snvs=[]
    all_ibd=[]
    ibd=str(1)
    i=0
    for snv in isec_set:
        loc,base,geno = snv.split('|')[0],snv.split('|')[1].replace('>','/'),snv.split('|')[2]
        chr,pos=loc.split('@')[0],loc.split('@')[1]
        chrN=chr.replace('chr','')
        if chrN=='X':
            chrN=23
        elif chrN=='Y':
            chrN=24
        elif chrN=='M':
            chrN=25
        
        all_snvs.append((chrN,pos,ibd,base,geno,0)) # This indixing is important, 0 means is shared for all samples
        all_ibd.append( (chrN,pos,ibd))
        #print ol, snv
        #print location_arr[i,:]
        i+=1
    print "The length of isec",len( isec_set.intersection(diff_snvs) )
    
    # SNVs not shared accross multiple individiduals.
    
    ibd=str(0)
    
    sp_dict=defaultdict()
    for ft, diff_snvs in snvs_singletons.iteritems():
        
        for snv in diff_snvs:
            loc,base,geno = snv.split('|')[0],snv.split('|')[1].replace('>','/'),snv.split('|')[2]
            chr,pos=loc.split('@')[0],loc.split('@')[1]        
            chrN=chr.replace('chr','')
    
            if chrN=='X':
                chrN=23
            elif chrN=='Y':
                chrN=24
            elif chrN=='M':
                chrN=25
            all_snvs.append((chrN,pos,ibd,base,geno,ft)) 
            all_ibd.append( (chrN,pos,ibd))              
            i+=1

    # Sort the snv table
    location_arr=np.array( all_snvs,dtype=mytype )
    location_arr_sorted = np.sort(location_arr, order=['CHRM', 'POS'])
    
    ########### Writing everything out
            
    # This takes time. Look for a better way to do it.    
    ibd_arr=np.array( all_ibd,dtype=mytype_ibd )
    ibd_arr_sorted=np.unique(ibd_arr)
    ibd_arr_sorted=np.sort(ibd_arr_sorted, order=['CHRM', 'POS'])
    
    of=open(options.ofile,'w')   
    of2=open(options.ofile+'_all','w')
    
    # Write all the snps.
    chrs = np.array(range(1,26))
    i=1
    for i in chrs:
        sub_array=location_arr_sorted[location_arr_sorted['CHRM']==i]
        of2=open(options.ofile+'_chr'+str(i)+'_all.txt','w')
        print sub_array
        for j in range(len(sub_array)):
            ol=list(sub_array[j])
            if ol[0]==23:
                ol[0]='X'
            elif ol[0]==24:
                ol[0]='Y'
            elif ol[0]==25: # CHRM M
                continue
            #print ",".join(map(str,ol)).replace(',','\t')
            of2.write(",".join(map(str,ol)).replace(',','\t')+'\n')
        of2.close()

'''
    # Write only the unique snps when ibd is considered
    for i in range(len(ibd_arr_sorted)):
        
        ol=list(ibd_arr_sorted[i])
        if ol[0]==23:
            ol[0]='X'
        elif ol[0]==24:
            ol[0]='Y'
        elif ol[0]==25: # CHRM M
            continue
        #print ",".join(map(str,ol)).replace(',','\t')
        of.write(",".join(map(str,ol)).replace(',','\t')+'\n')

    of.close()           
'''